#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
测试本地大语言模型是否正确集成到项目中
"""

import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from db.model_loader import model_loader
from config import API_TYPE


def test_local_llm():
    """测试本地大语言模型是否能正常工作"""
    print("=== 测试本地大语言模型集成 ===")
    
    if API_TYPE != "local":
        print(f"⚠️ 当前配置为使用 {API_TYPE} 模型，而非本地模型")
        print("请设置环境变量 API_TYPE=local 以使用本地模型")
        return False
    
    try:
        # 尝试加载本地LLM模型
        print("尝试加载本地大语言模型...")
        model_loader.init_llm()
        
        if not model_loader.get_llm() or not model_loader.get_llm_tokenizer():
            print("❌ 模型加载失败：无法获取有效的模型或分词器实例")
            return False
        
        print("✅ 模型加载成功！")
        
        # 测试生成简单回复
        test_prompt = "你好，请简单介绍一下自己。"
        print(f"\n测试提示词：{test_prompt}")
        
        # 构建完整提示，模拟chat.py中的格式
        full_prompt = f"你是一个智能助手，需要根据提供的文档内容和对话历史来回答用户的问题。请基于文档内容进行回答，不要添加文档中没有的信息。\n\n文档内容:\n测试文档内容\n\n对话历史:\nuser: {test_prompt}\nassistant:"
        
        print("生成回复中...（这可能需要一些时间）")
        response = model_loader.generate_llm_response(full_prompt)
        
        if response:
            print(f"✅ 回复生成成功！")
            print(f"回复内容：{response}")
            return True
        else:
            print("❌ 回复生成失败：返回了空内容")
            return False
    
    except Exception as e:
        print(f"❌ 测试过程中发生错误：{str(e)}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    success = test_local_llm()
    print("\n=== 测试总结 ===")
    if success:
        print("✅ 本地大语言模型集成测试通过！")
        sys.exit(0)
    else:
        print("❌ 本地大语言模型集成测试失败！")
        print("请检查以下几点：")
        print("1. 确认LLM模型已正确下载到配置的路径")
        print("2. 检查系统是否有足够的内存/显存来加载模型")
        print("3. 确保所有依赖都已正确安装")
        sys.exit(1)